#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient


# 连接MongoDB
def InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    client = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = client[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return client, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


# def RecurNode(tax_id, division_name, collection_taxonomy, collection_pak):
#     result = collection_taxonomy.find_one({"tax_id": tax_id})
#     if result["division_name"] == division_name:
#         if result["GenBank_hidden_flag"] == "0" or TestChildren(tax_id, division_name, collection_taxonomy) == True:
#             print("taxon:" + tax_id + "\t" + result["dms_name"] + "\t" + result["rank"])
#             collection_pak.update({"tax_id": result["tax_id"]}, {
#                 "$set": {"tax_id": result["tax_id"], "parent_tax_id": result["parent_tax_id"]["id"],
#                          "names_scientific_name": result["names_scientific_name"],
#                          "division_name": result["division_name"],
#                          "division_id": result["division_id"], "synonym": result["dms_synonym"],
#                          "rank": result["rank"]}},
#                                   upsert=True)
#             results = collection_taxonomy.find({"parent_tax_id.id": tax_id})
#             for item in results:
#                 RecurNode(item["tax_id"], division_name, collection_taxonomy, collection_pak)


def RecurNode(tax_id, division_name, collection_taxonomy, collection_pak):
    result = collection_taxonomy.find_one({"tax_id": tax_id})
    if result["division_name"] == division_name:
        print("taxon:" + tax_id + "\t" + result["dms_name"] + "\t" + result["rank"])
        collection_pak.update({"tax_id": result["tax_id"]}, {
            "$set": {"tax_id": result["tax_id"], "parent_tax_id": result["parent_tax_id"]["id"],
                     "names_scientific_name": result["names_scientific_name"],
                     "division_name": result["division_name"],
                     "division_id": result["division_id"], "synonym": result["dms_synonym"],
                     "rank": result["rank"]}},
                              upsert=True)
        if result["rank"] != "species":
            results = collection_taxonomy.find({"parent_tax_id.id": tax_id})
            for item in results:
                RecurNode(item["tax_id"], division_name, collection_taxonomy, collection_pak)


def Update_eLMSG(collection_elmsg, collection_pak):
    for result in collection_elmsg.find({"TaxonID": {"$exists": 1}}, {"TaxonID": 1, "LMSG_ID": 1}):
        collection_pak.update_one({"tax_id": result["TaxonID"]}, {"$set": {"LMSG_ID": result["LMSG_ID"]}})


# def TestChildren(tax_id, division_name, collection_taxonomy):
#     results = collection_taxonomy.find({"parent_tax_id.id": tax_id},
#                                        {"division_name": 1, "GenBank_hidden_flag": 1, "tax_id": 1})
#     for result in results:
#         if result["division_name"] == division_name:
#             if result["GenBank_hidden_flag"] == "0":
#                 return True
#             elif TestChildren(result["tax_id"], division_name, collection_taxonomy) == True:
#                 return True
#     return False


if __name__ == '__main__':
    client_taxonomy, collection_taxonomy = InitMongoDB("localhost", 27017, "Dictionary", "taxonomy_20181122")
    client_elmsg, collection_elmsg = InitMongoDB("localhost", 27017, "eLMSG", "eLMSG_20180612")
    client_pak, collection_pak = InitMongoDB("localhost", 27017, "PAK", "PAK")

    RecurNode("2", "Bacteria", collection_taxonomy, collection_pak)
    RecurNode("2157", "Bacteria", collection_taxonomy, collection_pak)
    # RecurNode("4751", "Plants and Fungi", collection_taxonomy, collection_pak)
    # RecurNode("10239", "Viruses", collection_taxonomy, collection_pak)

    Update_eLMSG(collection_elmsg, collection_pak)

    CloseMongoDB(client_taxonomy)
    CloseMongoDB(client_elmsg)
    CloseMongoDB(client_pak)
